#!/usr/bin/env python  
#-*- coding:utf-8 _*-  
""" 
@author:hello_life 
@license: Apache Licence 
@file: app.py 
@time: 2022/04/27
@software: PyCharm 
description:
"""
import torch
from transformers import AlbertTokenizer,AlbertModel
from flask import Flask,render_template,request

from parameters.albert_config import Config
from model.albert_model import Albert_model

app=Flask(__name__)

config=Config()
model=Albert_model(config)
tokenizer=AlbertTokenizer.from_pretrained(config.model_path)

@app.route('/')
def home():
    return render_template("page.html")

@app.route("/predict",methods=["POST"])
def predict():
    content=[i for i in request.form.values()]
    content="".join(content)
    content=tokenizer(content,padding="max_length",max_length=300,truncation=True,return_tensors="pt")
    model.load_state_dict(torch.load(config.save_model))
    with torch.no_grad():
        output=model(content)
        pred=output.argmax(1).item()
    return render_template("page.html",prediction_display_area=f"预测结果{pred}")



if __name__ == '__main__':
    app.run()